import segmentation_models_pytorch as smp
import torch
import cv2

img=cv2.imread("1.jpg")

aux_params=dict(
    pooling='avg',             # one of 'avg', 'max'
    dropout=0.5,               # dropout ratio, default is None
    activation='sigmoid',      # activatio、


    # n function, default is None
    classes=4,                 # define number of output labels
)

model = smp.DeepLabv3_Plus(
     encoder_name="resnet34",        # choose encoder, e.g. mobilenet_v2 or efficientnet-b7
     encoder_weights="imagenet",     # use `imagenet` pre-trained weights for encoder initialization
     in_channels=1,                  # model input channels (1 for gray-scale images, 3 for RGB, etc.)
     classes=3,                      # model output channels (number of classes in your dataset)
 )

preprocess_input = get_preprocessing_fn('resnet18', pretrained='imagenet')

model = smp.DeepLabV3Plus('resnet34', classes=4, aux_params=aux_params)
mask, label = model(torch.ones([64, 3, 512, 512]))
print(mask)

